import { Server as HTTPServer } from 'http';
import WebSocket from 'ws';
import { verifyToken } from './jwt';
import logger from './logger';
import os from 'os';

interface WebSocketClient extends WebSocket {
  userId?: string;
  tenantId?: string;
  isAlive?: boolean;
}

/**
 * WebSocket 服务器
 */
export class WebSocketServer {
  private wss: WebSocket.Server | null = null;
  private clients: Map<string, WebSocketClient> = new Map();
  private statusInterval: NodeJS.Timeout | null = null;

  /**
   * 初始化 WebSocket 服务器
   */
  initialize(server: HTTPServer) {
    this.wss = new WebSocket.Server({ 
      server,
      path: '/ws',
    });

    this.wss.on('connection', this.handleConnection.bind(this));

    // 启动心跳检测
    this.startHeartbeat();

    // 启动服务器状态推送
    this.startServerStatusBroadcast();

    logger.info('WebSocket 服务器已启动');
  }

  /**
   * 处理新连接
   */
  private async handleConnection(ws: WebSocketClient, req: any) {
    try {
      // 从查询参数中获取 token
      const url = new URL(req.url, `http://${req.headers.host}`);
      const token = url.searchParams.get('token');

      if (!token) {
        ws.close(1008, '未提供认证令牌');
        return;
      }

      // 验证 token
      let payload;
      try {
        payload = verifyToken(token);
      } catch (error) {
        ws.close(1008, '无效的认证令牌');
        return;
      }

      // 保存用户信息
      ws.userId = payload.userId;
      ws.tenantId = payload.tenantId || undefined;
      ws.isAlive = true;

      // 添加到客户端列表
      this.clients.set(payload.userId, ws);

      logger.info(`WebSocket 客户端已连接: ${payload.userId}`);

      // 发送欢迎消息
      this.sendMessage(ws, {
        type: 'connected',
        message: '已连接到服务器',
        timestamp: Date.now(),
      });

      // 监听消息
      ws.on('message', (data: string) => {
        this.handleMessage(ws, data);
      });

      // 监听心跳响应
      ws.on('pong', () => {
        ws.isAlive = true;
      });

      // 监听关闭
      ws.on('close', () => {
        if (ws.userId) {
          this.clients.delete(ws.userId);
          logger.info(`WebSocket 客户端已断开: ${ws.userId}`);
        }
      });

      // 监听错误
      ws.on('error', (error) => {
        logger.error(`WebSocket 错误: ${error.message}`);
      });
    } catch (error) {
      logger.error('WebSocket 连接处理失败:', error);
      ws.close(1011, '服务器内部错误');
    }
  }

  /**
   * 处理客户端消息
   */
  private handleMessage(ws: WebSocketClient, data: string) {
    try {
      const message = JSON.parse(data.toString());
      
      switch (message.type) {
        case 'ping':
          this.sendMessage(ws, { type: 'pong', timestamp: Date.now() });
          break;
        case 'subscribe':
          // 处理订阅请求
          this.handleSubscribe(ws, message.channel);
          break;
        default:
          logger.warn(`未知的消息类型: ${message.type}`);
      }
    } catch (error) {
      logger.error('处理 WebSocket 消息失败:', error);
    }
  }

  /**
   * 处理订阅请求
   */
  private handleSubscribe(ws: WebSocketClient, channel: string) {
    logger.info(`客户端订阅频道: ${channel}`);
    this.sendMessage(ws, {
      type: 'subscribed',
      channel,
      timestamp: Date.now(),
    });
  }

  /**
   * 发送消息给指定客户端
   */
  private sendMessage(ws: WebSocket, data: any) {
    if (ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify(data));
    }
  }

  /**
   * 广播消息给所有客户端
   */
  broadcast(data: any, filter?: (client: WebSocketClient) => boolean) {
    this.clients.forEach((client) => {
      if (client.readyState === WebSocket.OPEN) {
        if (!filter || filter(client)) {
          this.sendMessage(client, data);
        }
      }
    });
  }

  /**
   * 发送消息给指定用户
   */
  sendToUser(userId: string, data: any) {
    const client = this.clients.get(userId);
    if (client) {
      this.sendMessage(client, data);
    }
  }

  /**
   * 发送消息给指定租户的所有用户
   */
  sendToTenant(tenantId: string, data: any) {
    this.broadcast(data, (client) => client.tenantId === tenantId);
  }

  /**
   * 启动心跳检测
   */
  private startHeartbeat() {
    setInterval(() => {
      this.clients.forEach((client, userId) => {
        if (client.isAlive === false) {
          this.clients.delete(userId);
          return client.terminate();
        }

        client.isAlive = false;
        client.ping();
      });
    }, 30000); // 每30秒检测一次
  }

  /**
   * 启动服务器状态推送
   */
  private startServerStatusBroadcast() {
    this.statusInterval = setInterval(() => {
      const status = this.getServerStatus();
      this.broadcast({
        type: 'server_status',
        data: status,
        timestamp: Date.now(),
      });
    }, 5000); // 每5秒推送一次
  }

  /**
   * 获取服务器状态
   */
  private getServerStatus() {
    const cpuUsage = process.cpuUsage();
    const memUsage = process.memoryUsage();

    return {
      cpu: {
        user: cpuUsage.user,
        system: cpuUsage.system,
      },
      memory: {
        rss: (memUsage.rss / 1024 / 1024).toFixed(2),
        heapTotal: (memUsage.heapTotal / 1024 / 1024).toFixed(2),
        heapUsed: (memUsage.heapUsed / 1024 / 1024).toFixed(2),
        heapUsedPercent: ((memUsage.heapUsed / memUsage.heapTotal) * 100).toFixed(2),
      },
      system: {
        totalMemory: (os.totalmem() / 1024 / 1024 / 1024).toFixed(2),
        freeMemory: (os.freemem() / 1024 / 1024 / 1024).toFixed(2),
        usedMemory: ((os.totalmem() - os.freemem()) / 1024 / 1024 / 1024).toFixed(2),
        memoryUsagePercent: (((os.totalmem() - os.freemem()) / os.totalmem()) * 100).toFixed(2),
        loadavg: os.loadavg(),
        cpuCount: os.cpus().length,
      },
    };
  }

  /**
   * 关闭 WebSocket 服务器
   */
  close() {
    if (this.statusInterval) {
      clearInterval(this.statusInterval);
    }

    if (this.wss) {
      this.wss.close(() => {
        logger.info('WebSocket 服务器已关闭');
      });
    }
  }

  /**
   * 获取连接数
   */
  getConnectionCount() {
    return this.clients.size;
  }
}

export const wsServer = new WebSocketServer();

